AWS Cloud Development Kit(CDK)(Python)を使って「VPC+Sagemaker+Lambda」環境構築してみた(Lambda編)

suto

2020.08.18

この記事は公開されてから1年以上経過しています。情報が古い可能性がありますので、ご注意ください。

どーもsutoです。

前回の記事でAWS Cloud Development Kit(CDK)を使ってPythonコードでSagemakerノートブックインスタンス構築まで行ってきました。

前々回のVPC編はこちらです。

今回はsagemakerのノートブックインスタンスをスケジュール停止するLambda関数を追加していきます。

はじめに

今回、自分のアカウントで機械学習検証用環境のインフラを構築し、いつでも環境のデプロイ／破棄ができるように自動化＆コード管理したい、というモチベからAWS CDKを本格的に触ってみることにしました。

本検証の目標

AWS CDKのインストールと初期設定、VPC構築
Sagemakerのノートブックインスタンスを自動構築、VPCとの関連付け
ノートブックインスタンスのスケジュールによる自動停止を行うLambdaの追加設定（←本記事）

でやっていきます。

CDK環境にSagemakerモジュール追加

前回作成したsage-nwプロジェクトの仮想環境ディレクトリ内に入っている状態からスタートします。

 
~ sage-nw % source .env/bin/activate
(.env) ~ sage-nw %

setup.pyの「install_requires」部分を編集して必要なモジュールをインストールします。

 
    install_requires=[
        "aws-cdk.core",
        "aws-cdk.aws-ec2",
        "aws-cdk.aws-sagemaker",
        "aws_cdk.aws_events",
        "aws_cdk.aws_lambda",
        "boto3",
        ],

 
pip install -r requirements.txt

Lambda関数スクリプトをフォルダに格納

今回のSagemaker構築内容は以下のとおりです。

指定のタグがついたノートブックインスタンスをスケジュールで自動停止するLambda関数作成

以下の記事のLambda関数を使わせていただき、これをAWS CDKからデプロイさせるようにします。（Lambdaの詳しい処理内容はリンク先の記事を参照ください）

まず、「sage-nw」フォルダで「lambda」フォルダを作成し、そのなかに上記記事のLambda関数を格納します。


~ sage-nw % mkdir lambda && cd lambda

autostop.py

 
import json
import boto3
import os

# 環境変数を読み込む

TAG_KEY = os.environ['TAG_KEY']
TAG_VALUE = os.environ['TAG_VALUE']


def handler(event, context):

    # SageMakerが対応してるリージョン一覧を取得
    #regions = boto3.Session().get_available_regions('sagemaker')
    regions = ['ap-northeast-1', 'ap-northeast-2', 'ap-south-1', 'ap-southeast-1', 'ap-southeast-2', 'ca-central-1', 'eu-central-1', 'eu-west-1', 'eu-west-2', 'eu-west-3', 'sa-east-1', 'us-east-1', 'us-east-2', 'us-west-1', 'us-west-2']
    
    # リージョンごとに実行
    for region in regions:
        sm = boto3.client('sagemaker', region_name=region)
    
        # 起動中(InService)のインスタンスを取得する
        nb_list = sm.list_notebook_instances(StatusEquals='InService')['NotebookInstances']
    
        for nb in nb_list:
            # インスタンスのタグを取得する
            tags = sm.list_tags(ResourceArn=nb['NotebookInstanceArn'])['Tags']
    
            # タグから自動停止が有効かどうかを調べる
            stop_nb = False
            for tag in tags:
                if tag['Key'] == TAG_KEY and tag['Value'] == TAG_VALUE:
                    stop_nb = True
                    break
    
            # 自動停止が有効なノートブックインスタンスであれば停止させる
            if stop_nb:
                nb_name = nb['NotebookInstanceName']
                print('stop', nb_name+'@'+region)
                stop_notebook_instance(nb_name, sm)
    
    return {
        'statusCode': 200,
        'body': json.dumps('done')
    }



def stop_notebook_instance(nb_name, sm_client):
    # 指定されたノートブックインスタンスを停止する

    sm_client.stop_notebook_instance(NotebookInstanceName=nb_name)

[小ネタ]Lambda関数の動作テストでハマった話

当初Sagemaker対応リージョンからインスタンス一覧を取得する際、boto3.Session().get_available_regions('sagemaker')を使用したのですが、エラー発生で取得できませんでした。

原因は私のアカウントで一部のAWSリージョンが有効化していなかったためでした。現在、香港やバーレーンなどの一部リージョンはデフォルトで無効になっているため、上記コマンドのような全リージョンから情報取得する際は注意が必要です。

有効リージョンの確認と有効化/無効化手順

今回は応急処置として、情報を取得するリージョンを選定してリストに突っ込みました。（使わないリージョンをわざわざ有効化しなくてもよいかなーと思いまして）

コードの記述

「sage_nw」フォルダに移動してlambda_stack.pyを新規作成し、以下のコードを記述します。


from aws_cdk import  (
    core,
    aws_lambda as _lambda,
    aws_events as events,
    aws_events_targets as targets,
)

class LambdaStack(core.Construct):

    def __init__(self, scope: core.Construct, id: str, **kwargs) -> None:
        super().__init__(scope, id, **kwargs)

        # The code that defines your stack goes here

        # Defines an AWS Lambda resource
        my_lambda = _lambda.Function(
            self, 'AutoStopHandler',
            runtime=_lambda.Runtime.PYTHON_3_7,
            code=_lambda.Code.asset('lambda'),
            handler='autostop.handler',
            environment={
                'TAG_KEY': "auto_stop",
                'TAG_VALUE': "true",
            },
        )
        # LambdaをトリガーするEventBridge作成
        events_sutostop = events.Rule(
            self,
            "ScheduleRule",
            schedule=events.Schedule.cron(
                minute='0',
                hour='18',
                month='*',
                week_day='MON-FRI',
                year='*'),
        )
        events_sutostop.add_target(targets.LambdaFunction(my_lambda))

sage_nw_stack.pyにlambda_stackの処理を呼び出す内容を追記して保存します。

from aws_cdk import core
from aws_cdk import aws_ec2 as ec2
from sagemaker_stack import SagemakerStack
from lambda_stack import LambdaStack

class SageNwStack(core.Stack):

    def __init__(self, scope: core.Construct, id: str, **kwargs) -> None:
        super().__init__(scope, id, **kwargs)

        # The code that defines your stack goes here
        # 変数の宣言
        vpc_cidr = '10.1.0.0/16'
        subnet_mask = 24

        # 新規VPC作成
        vpc = ec2.Vpc(
            self,
            id="Sage-vpc",
            cidr=vpc_cidr,
            nat_gateways=0,
            subnet_configuration=[
                ec2.SubnetConfiguration(
                    cidr_mask=subnet_mask, name='public', subnet_type=ec2.SubnetType.PUBLIC,
                ),
                ec2.SubnetConfiguration(
                    cidr_mask=subnet_mask, name='private', subnet_type=ec2.SubnetType.ISOLATED,
                ),
            ],
        )

        security_group = ec2.SecurityGroup(
            self,
            id='Sage-sg',
            vpc=vpc,
            security_group_name='test-sg'
        )
        security_group.add_ingress_rule(
            peer=ec2.Peer.ipv4(vpc_cidr),
            connection=ec2.Port.all_traffic(),
        )

        # vpcメソッドの結果を変数に格納
        subnet = vpc.isolated_subnets[0].subnet_id
        sg = [security_group.security_group_id]

        # ノートブックインスタンス作成
        sagemaker = SagemakerStack(
            self,
            "Notebook",
            subnetid = subnet,
            sgid = sg,
        )
        # Lambda作成
        lambda_autostop = LambdaStack(
            self,
            "AutoStop",
        )

VPC,Sagemaker編を含めた最終的なディレクトリ構成は以下となります。

スタックの確認とデプロイ

cdk diffコマンドでデプロイ済の Stack との差異を表示して内容を確認します。

cdk diff sage-nw --profile suto

確認が終わったら以下コマンドでスタックを更新してみます。

 
(.env) ~ sage-nw % cdk deploy sage-nw --profile suto

スタックが更新されていることを確認できました。

まとめ

よくあるLambda＋EventBridgeの設定を追加し、環境としては全体的に大きなボリュームになりましたが、機能ごとにコードを細分化し、ステップを踏んで作成していけばけっこうわかりやすくなったかなと思います。

AWS CDKはTypescriptで記述したブログは多く見かけますが、Pythonの記事は少ないと感じていたため、ブログに起こすことでPython使いのみなさまに少しでもお役に立てればと思っています。

AWS Cloud Development Kit(CDK)(Python)を使って「VPC+Sagemaker+Lambda」環境構築してみた(Lambda編)

はじめに

本検証の目標

CDK環境にSagemakerモジュール追加

Lambda関数スクリプトをフォルダに格納

[小ネタ]Lambda関数の動作テストでハマった話

コードの記述

スタックの確認とデプロイ

まとめ

関連記事

主なカテゴリ

AWS

おすすめ

プロダクト

コンテンツ

お問い合わせ

運営会社